
#include "cudaBindings.h"
#include "../Common/InteractionLogic.h"

//Allocates memory on gpu and copies required data.
inline float *copyToGpu(float *data, size_t count)
{
	size_t size = sizeof(float) * count;
	float *res;
	SAFE_CALL(cudaMalloc((void **)&res, size));
	SAFE_CALL(cudaMemcpy(res, data, size, cudaMemcpyHostToDevice));
	return res;
}

//Copies data from GPU and releases allocated GPU-memory.
inline void copyFromGpu(float *data_host, float *&data_dev, size_t count)
{
	size_t size = sizeof(float) * count;
	SAFE_CALL(cudaMemcpy(data_host, data_dev, size, cudaMemcpyDeviceToHost));
	SAFE_CALL(cudaFree(data_dev));
	data_dev = NULL;
}

//---------------------------
//--- cudaBodyInteractor1 ---
//---------------------------

//Interacts all bodies of required range.
__global__ void interactRangeKernel(float *px, float *py, float *pz,
									float *ax, float *ay, float *az,
									float *masses, size_t count)
{
	int n = blockIdx.x * blockDim.x + threadIdx.x;
	if (n < count)
	{
		//Storing body's parameters in registers.
		float _px = px[n];
		float _py = py[n];
		float _pz = pz[n];
		float _ax = ax[n];
		float _ay = ay[n];
		float _az = az[n];
		float _mass = masses[n];
		//Interacting.
		for (int i = 0; i < count; i++)
		{
			if (i != n)
				InteractionLogic::interactWith(_px, _py, _pz,
											   _ax, _ay, _az,
											   _mass,
											   px[i], py[i], pz[i],
											   masses[i]);
		}
		//Copying acceleration to memory.
		ax[n] = _ax;
		ay[n] = _ay;
		az[n] = _az;
	}
}

void cudaBodyInteractor1::interactRange(InterleavedBodies *bodies, size_t lo, size_t hi)
{
	//Copying data to GPU memory.
	size_t count = hi - lo;
	float *px_dev = copyToGpu(bodies->positions.getX() + lo, count);
	float *py_dev = copyToGpu(bodies->positions.getY() + lo, count);
	float *pz_dev = copyToGpu(bodies->positions.getZ() + lo, count);
	float *ax_dev = copyToGpu(bodies->accelerations.getX() + lo, count);
	float *ay_dev = copyToGpu(bodies->accelerations.getY() + lo, count);
	float *az_dev = copyToGpu(bodies->accelerations.getZ() + lo, count);
	float *masses_dev = copyToGpu(&bodies->masses[0] + lo, count);
	//Launching kernel.
	dim3 threads(512);
	dim3 grid(count / 512 + 1);
	interactRangeKernel<<<threads, grid>>>(px_dev, py_dev, pz_dev,
										   ax_dev, ay_dev, az_dev,
										   masses_dev, count);
	CHECK_ERROR("\"interactRangeKernel\" kernel failed");
	//Copying results and freeing GPU memory.
	SAFE_CALL(cudaFree(px_dev));
	SAFE_CALL(cudaFree(py_dev));
	SAFE_CALL(cudaFree(pz_dev));
	SAFE_CALL(cudaFree(masses_dev));
	copyFromGpu(bodies->accelerations.getX() + lo, ax_dev, count);
	copyFromGpu(bodies->accelerations.getY() + lo, ay_dev, count);
	copyFromGpu(bodies->accelerations.getZ() + lo, az_dev, count);
}

//---------------------------
//--- cudaBodyInteractor2 ---
//---------------------------

__global__ void interactRangesKernel(float *px1, float *py1, float *pz1,
									 float *ax1, float *ay1, float *az1,
									 float *masses1, size_t count1,
									 float *px2, float *py2, float *pz2,
									 float *ax2, float *ay2, float *az2,
									 float *masses2, size_t count2)
{
	int n = blockIdx.x * blockDim.x + threadIdx.x;
	if (n < count1)
	{
		//Storing body's parameters in registers.
		float _px1 = px1[n];
		float _py1 = py1[n];
		float _pz1 = pz1[n];
		float _ax1 = ax1[n];
		float _ay1 = ay1[n];
		float _az1 = az1[n];
		float _mass1 = masses1[n];
		//Interacting.
		for (int i = 0; i < count2; i++)
		{
			int j = (i + n) % count2;
			float _ax2 = ax2[j];
			float _ay2 = ay2[j];
			float _az2 = az2[j];
			InteractionLogic::interactWithEachOther(_px1, _py1, _pz1,
													_ax1, _ay1, _az1,
													_mass1,
													px2[j], py2[j], pz2[j],
													_ax2, _ay2, _az2,
													masses2[j]);
			ax2[j] = _ax2;
			ay2[j] = _ay2;
			az2[j] = _az2;
			__syncthreads();
		}
		//Copying acceleration to memory.
		ax1[n] = _ax1;
		ay1[n] = _ay1;
		az1[n] = _az1;
	}
}

void cudaBodyInteractor2::interactRanges(InterleavedBodies *b1, size_t lo1, size_t hi1,
										 InterleavedBodies *b2, size_t lo2, size_t hi2)
{
	//Copying data to GPU memory.
	size_t count1 = hi1 - lo1;
	size_t count2 = hi2 - lo2;
		//First range.
	float *px1_dev = copyToGpu(b1->positions.getX() + lo1, count1);
	float *py1_dev = copyToGpu(b1->positions.getY() + lo1, count1);
	float *pz1_dev = copyToGpu(b1->positions.getZ() + lo1, count1);
	float *ax1_dev = copyToGpu(b1->accelerations.getX() + lo1, count1);
	float *ay1_dev = copyToGpu(b1->accelerations.getY() + lo1, count1);
	float *az1_dev = copyToGpu(b1->accelerations.getZ() + lo1, count1);
	float *masses1_dev = copyToGpu(&b1->masses[0] + lo1, count1);
		//Second range.
	float *px2_dev = copyToGpu(b2->positions.getX() + lo2, count2);
	float *py2_dev = copyToGpu(b2->positions.getY() + lo2, count2);
	float *pz2_dev = copyToGpu(b2->positions.getZ() + lo2, count2);
	float *ax2_dev = copyToGpu(b2->accelerations.getX() + lo2, count2);
	float *ay2_dev = copyToGpu(b2->accelerations.getY() + lo2, count2);
	float *az2_dev = copyToGpu(b2->accelerations.getZ() + lo2, count2);
	float *masses2_dev = copyToGpu(&b2->masses[0] + lo2, count2);
	//Launching kernel.
	dim3 threads(512);
	dim3 grid(count1 / 512 + 1);
	interactRangesKernel<<<threads, grid>>>(px1_dev, py1_dev, pz1_dev,
										    ax1_dev, ay1_dev, az1_dev,
										    masses1_dev, count1,
										    px2_dev, py2_dev, pz2_dev,
										    ax2_dev, ay2_dev, az2_dev,
										    masses2_dev, count2);
	CHECK_ERROR("\"interactRangeKernel\" kernel failed");
	//Copying results and freeing GPU memory.
	SAFE_CALL(cudaFree(px1_dev));
	SAFE_CALL(cudaFree(py1_dev));
	SAFE_CALL(cudaFree(pz1_dev));
	SAFE_CALL(cudaFree(px2_dev));
	SAFE_CALL(cudaFree(py2_dev));
	SAFE_CALL(cudaFree(pz2_dev));
	SAFE_CALL(cudaFree(masses1_dev));
	SAFE_CALL(cudaFree(masses2_dev));
	copyFromGpu(b1->accelerations.getX() + lo1, ax1_dev, count1);
	copyFromGpu(b1->accelerations.getY() + lo1, ay1_dev, count1);
	copyFromGpu(b1->accelerations.getZ() + lo1, az1_dev, count1);
	copyFromGpu(b2->accelerations.getX() + lo2, ax2_dev, count2);
	copyFromGpu(b2->accelerations.getY() + lo2, ay2_dev, count2);
	copyFromGpu(b2->accelerations.getZ() + lo2, az2_dev, count2);
}